"""
https://blog.csdn.net/sugarannie/article/details/53080168
OpenCV库中watershed函数（分水岭算法）的详细使用例程
【c++转python】
【GUI貌似不支持中文，窗口标题替换成英文】
"""

import cv2 as cv
import numpy as np
import os

np.random.seed(1)


# 操作提示
def helpText():
    help = """
    先用鼠标在图片窗口中标记出大致的区域
	如果想把图片分割为N个区域，就要做N个标记
	键盘按键【1】	- 运行的分水岭分割算法
	键盘按键【2】	- 恢复原始图片
	键盘按键【0】	- 依次分割每个区域（必须先按【1】）
	键盘按键【ESC】	- 退出程序
    """
    print(help)


def on_Mouse(event, x, y, flags, param):
    global clickPoint

    # 如果鼠标不在窗口中则返回
    if x < 0 or x >= srcImage.shape[1]\
        or y < 0 or y >= srcImage.shape[0]:
        return

    # 如果鼠标左键被按下，获取鼠标当前位置；当鼠标左键按下并且移动时，绘制白线；
    if event == cv.EVENT_LBUTTONDOWN:
        clickPoint = x, y
    elif event == cv.EVENT_MOUSEMOVE and flags & cv.EVENT_FLAG_LBUTTON:
        point = (x, y)
        cv.line(maskImage, clickPoint, point, 255, 5, cv.LINE_8)
        cv.line(srcImage, clickPoint, point, (255, 255, 255), 5, cv.LINE_8)
        clickPoint = point
        cv.imshow(title, srcImage)


helpText()  # 操作提示

# dir = '../../../../../large_data/pic/watershed/'
# name = 'air.jpg'
dir = '../../../../../large_data/pic/'
name = 'dog_bird.png'
path = os.path.join(dir, name)
srcImage = cv.imread(path)
# srcImage = cv.imread(path, cv.IMREAD_COLOR)
srcImage_ = srcImage.copy()  # 程序中srcImage会被改变，所以这里做备份
maskImage = np.zeros(srcImage.shape[:2], dtype=np.uint8)  # 掩模，在上面做标记，然后传给findContours
clickPoint = (0, 0)  # 鼠标点下去的位置

areaCount = 1  # 计数，在按【0】时绘制每个区域

title = 'Mark in pic'
cv.imshow(title, srcImage)

cv.setMouseCallback(title, on_Mouse)

while True:
    c = cv.waitKey(0) & 0xFF

    if 27 == c:  # 按【ESC】键退出
        break

    if ord('2') == c:  # 按【2】恢复原图
        maskImage = np.zeros(srcImage.shape[:2], dtype=np.uint8)
        srcImage = srcImage_.copy()
        cv.imshow(title, srcImage)

    if ord('1') == c:  # 按【1】处理图片
        contours, _ = cv.findContours(maskImage, cv.RETR_CCOMP, cv.CHAIN_APPROX_SIMPLE)

        if len(contours) == 0:  # 如果没有做标记，即没有轮廓，则退出该if语句
            break
        print(f'{len(contours)}个轮廓')

        maskWaterShed = np.zeros(maskImage.shape, dtype=np.int32)  # watershed()函数的参数

        # 【注意】 将颜色表比原代码提前定义，以便在做分水岭之前就检查maskWaterShed
        # 【注意】 将颜色表比原代码多定义索引0为白色，以便利用numpy强大的切片代替c++里的多层循环
        # color table start
        # 随机生成几种颜色
        colorTab = [(255, 255, 255)]
        for i in range(len(contours)):
            b = np.random.randint(0, 256)
            g = np.random.randint(0, 256)
            r = np.random.randint(0, 256)
            colorTab.append((b, g, r))
        colorTab = np.uint8(colorTab)
        # color table end

        # 在maskWaterShed上绘制轮廓
        for i in range(len(contours)):
            cv.drawContours(maskWaterShed, contours, i, i + 1, -1)

        print('分水岭前: ', np.unique(maskWaterShed))  # 【注意】maskWaterShed中的元素
        resImage = colorTab[maskWaterShed]  # 【注意】利用numpy强大的切片代替c++里的多层循环
        cv.imshow('before watershed', resImage)  # 【注意】在做分水岭之前就显示maskWaterShed

        # 分水岭
        cv.watershed(srcImage_, maskWaterShed)
        print('分水岭后: ', np.unique(maskWaterShed))  # 【注意】maskWaterShed中的元素

        maskWaterShed[maskWaterShed == -1] = 0
        print('分水岭效果显示前: ', np.unique(maskWaterShed))  # 【注意】maskWaterShed中的元素

        # 最后要显示的图像
        resImage = colorTab[maskWaterShed]  # 【注意】利用numpy强大的切片代替c++里的多层循环

        cv.imshow('resImage', resImage)
        resImage = cv.addWeighted(resImage, 0.3, srcImage_, 0.7, 0)
        cv.imshow("Watershed result", resImage)

    if ord('0') == c:  # 多次点按【0】依次显示每个被分割的区域，需要先按【1】处理图像
        resImage = np.zeros((*maskWaterShed.shape, 3), dtype=np.uint8)  # 声明一个最后要显示的图像
        resImage[maskWaterShed == areaCount] = srcImage_[maskWaterShed == areaCount]  # 【注意】利用numpy强大的切片代替c++里的多层循环
        resImage[maskWaterShed != areaCount] = (0, 0, 0)  # 【注意】利用numpy强大的切片代替c++里的多层循环
        cv.imshow("Watershed result", resImage)
        areaCount += 1
        if areaCount >= 4:
            areaCount = 1

print('Existed')
cv.destroyAllWindows()
print('Windows closed')
